import os
import json
from sklearn.metrics import confusion_matrix
import torch
from PIL import Image
from torchvision import transforms
import numpy as np
from ResNet import resnet50
import matplotlib.pyplot as plt


classarr = ['Center', 'Donut', 'Edge-Loc', 'Edge-Ring', 'Loc', 'Near-full', 'Random', 'Scratch']
def main():
    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")

    data_transform = transforms.Compose([
        transforms.Resize((26, 26)),
        transforms.ToTensor(),
        transforms.Normalize(mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5))])

    # load image
    # 指向需要遍历预测的图像文件夹
    class_names = ['Center', 'Donut', 'Edge-Loc', 'Edge-Ring', 'Loc', 'Near-full', 'Random', 'Scratch']
    img_path_list = []
    for i in range(len(class_names)):
        imgs_root = "C://Users//w8887757//PycharmProjects//wm811k//data//test//" + class_names[i]
        img_path_list += [os.path.join(imgs_root, i) for i in os.listdir(imgs_root) if i.endswith(".jpg")]

    # imgs_root = "./input/val"
    # assert os.path.exists(imgs_root), f"file: '{imgs_root}' dose not exist."
    # # 读取指定文件夹下所有jpg图像路径
    # img_path_list = [os.path.join(imgs_root, i) for i in os.listdir(imgs_root) if i.endswith(".jpg")]

    # read class_indict
    class_indict = {'0':'Center','1':'Donut','2':'Edge-Loc','3':'Edge-Ring','4':'Loc','5':'Near-full','6':'Random','7':'Scratch'}

    # create model
    model = resnet50(num_classes=8).to(device)

    # load model weights
    weights_path = "./input/save_weights/resnet50_dataAug_epo150_class8.pth"
    assert os.path.exists(weights_path), f"file: '{weights_path}' dose not exist."
    model.load_state_dict(torch.load(weights_path, map_location=device))

    # prediction
    model.eval()
    batch_size = 8  # 每次预测时将多少张图片打包成一个batch
    y_true = []
    y_pred = []
    truecnt = 0;
    with torch.no_grad():
        for ids in range(0, len(img_path_list) // batch_size):
            img_list = []
            for img_path in img_path_list[ids * batch_size: (ids + 1) * batch_size]:
                assert os.path.exists(img_path), f"file: '{img_path}' dose not exist."
                img = Image.open(img_path)
                img = data_transform(img)
                img_list.append(img)

            # batch img
            # 将img_list列表中的所有图像打包成一个batch
            batch_img = torch.stack(img_list, dim=0)
            # predict class
            output = model(batch_img.to(device)).cpu()
            predict = torch.softmax(output, dim=1)
            probs, classes = torch.max(predict, dim=1)

            for idx, (pro, cla) in enumerate(zip(probs, classes)):
                print(str(cla.numpy()))
                y_pred.append(class_indict[str(cla.numpy())])
                y_true.append(str(img_path_list[ids * batch_size + idx])[len("C://Users//w8887757//PycharmProjects//wm811k//data//test//"):].split("\\")[0])
                if (class_indict[str(cla.numpy())] == str(img_path_list[ids * batch_size + idx])[len("C://Users//w8887757//PycharmProjects//wm811k//data//test//"):].split("\\")[0]):
                    truecnt = truecnt + 1
                print("image: {}  class: {}  prob: {:.3}".format(img_path_list[ids * batch_size + idx],
                                                                 class_indict[str(cla.numpy())],
                                                                 pro.numpy()))
    predict_acc = truecnt /len(y_true)
    print("准确率：" + str(predict_acc))
    cm = confusion_matrix(y_true, y_pred)
    plot_confusion_matrix(cm, 'confusion_matrix_data8_epo100.png', title='confusion matrix')

def plot_confusion_matrix(cm, savename, title='Confusion Matrix'):
    plt.figure(figsize=(12, 8), dpi=100)
    np.set_printoptions(precision=2)

    # 在混淆矩阵中每格的概率值
    ind_array = np.arange(len(classarr))
    x, y = np.meshgrid(ind_array, ind_array)
    for x_val, y_val in zip(x.flatten(), y.flatten()):
        c = cm[y_val][x_val]
        if c > 0.001:
            plt.text(x_val, y_val, "%d" % (c,), color='red', fontsize=15, va='center', ha='center')

    plt.imshow(cm, interpolation='nearest', cmap=plt.cm.binary)
    plt.title(title)
    plt.colorbar()
    xlocations = np.array(range(len(classarr)))
    plt.xticks(xlocations, classarr, rotation=90)
    plt.yticks(xlocations, classarr)
    plt.ylabel('Actual label')
    plt.xlabel('Predict label')

    # offset the tick
    tick_marks = np.array(range(len(classarr))) + 0.5
    plt.gca().set_xticks(tick_marks, minor=True)
    plt.gca().set_yticks(tick_marks, minor=True)
    plt.gca().xaxis.set_ticks_position('none')
    plt.gca().yaxis.set_ticks_position('none')
    plt.grid(True, which='minor', linestyle='-')
    plt.gcf().subplots_adjust(bottom=0.15)

    # show confusion matrix
    plt.savefig(savename, format='png')
    plt.show()


if __name__ == '__main__':
    main()